package com.xkcyy.one.ai.common.tool.util;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import com.github.victools.jsonschema.module.jackson.JacksonModule;
import com.github.victools.jsonschema.module.jackson.JacksonOption;
import com.github.victools.jsonschema.module.swagger2.Swagger2Module;
import com.xkcyy.one.ai.common.annotation.Tool;
import com.xkcyy.one.ai.common.annotation.ToolParam;
import com.xkcyy.one.ai.llm.model.ToolDefinition;
import com.xkcyy.one.ai.llm.model.ToolDefinitionFunction;
import com.xkcyy.one.ai.llm.model.ToolDefinitionParameters;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.media.Schema.RequiredMode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.stream.Stream;

/**
 * JSON Schema生成器
 * 
 * @author yuand
 */
@Slf4j
public final class JsonSchemaGenerator {
    private static final boolean PROPERTY_REQUIRED_BY_DEFAULT = true;
    private static final SchemaGenerator TYPE_SCHEMA_GENERATOR;
    private static final SchemaGenerator SUBTYPE_SCHEMA_GENERATOR;

    private JsonSchemaGenerator() {
    }

    /**
     * 生成工具定义信息
     */
    public static ToolDefinition generateToolDefinition(Class<?> toolClass) {
        Tool toolAnnotation = toolClass.getAnnotation(Tool.class);
        if (toolAnnotation == null) {
            throw new IllegalArgumentException("Class must be annotated with @Tool");
        }

        // 创建工具函数定义
        ToolDefinitionFunction function = ToolDefinitionFunction.builder()
                .name(toolAnnotation.name())
                .description(toolAnnotation.description())
                .build();

        // 获取所有公共方法
        Method[] methods = toolClass.getMethods();
        for (Method method : methods) {
            if (method.getDeclaringClass() == Object.class) {
                continue;
            }

            try {
                // 生成方法参数的Schema
                String inputSchema = generateForMethodInput(new ObjectMapper(), method);
                ToolDefinitionParameters parameters = new ObjectMapper().readValue(inputSchema,
                        ToolDefinitionParameters.class);
                function.setParameters(parameters);
            } catch (JsonProcessingException e) {
                log.error("Failed to parse method parameters schema", e);
                throw new RuntimeException("Failed to parse method parameters schema", e);
            }
        }

        // 创建工具定义
        return ToolDefinition.builder()
                .type("function")
                .function(function)
                .build();
    }

    /**
     * 生成方法参数的Schema
     */
    public static String generateForMethodInput(ObjectMapper objectMapper, Method method) {
        ObjectNode schema = objectMapper.createObjectNode();
        schema.put("$schema", SchemaVersion.DRAFT_2020_12.getIdentifier());
        schema.put("type", "object");
        ObjectNode properties = schema.putObject("properties");
        List<String> required = new ArrayList<>();

        for (int i = 0; i < method.getParameterCount(); ++i) {
            String parameterName = method.getParameters()[i].getName();
            Type parameterType = method.getGenericParameterTypes()[i];
            if (isMethodParameterRequired(method, i)) {
                required.add(parameterName);
            }

            ObjectNode parameterNode = SUBTYPE_SCHEMA_GENERATOR.generateSchema(parameterType, new Type[0]);
            String parameterDescription = getMethodParameterDescription(method, i);
            if (StringUtils.hasText(parameterDescription)) {
                parameterNode.put("description", parameterDescription);
            }

            properties.set(parameterName, parameterNode);
        }

        ArrayNode requiredArray = schema.putArray("required");
        Objects.requireNonNull(requiredArray);
        required.forEach(requiredArray::add);
        return schema.toPrettyString();
    }

    /**
     * 生成类型Schema
     */
    public static String generateForType(Type type, SchemaOption... schemaOptions) {
        Assert.notNull(type, "type cannot be null");
        ObjectNode schema = TYPE_SCHEMA_GENERATOR.generateSchema(type, new Type[0]);
        if (type == Void.class && !schema.has("properties")) {
            schema.putObject("properties");
        }

        processSchemaOptions(schemaOptions, schema);
        return schema.toPrettyString();
    }

    private static void processSchemaOptions(SchemaOption[] schemaOptions, ObjectNode schema) {
        if (Stream.of(schemaOptions)
                .noneMatch((option) -> option == SchemaOption.ALLOW_ADDITIONAL_PROPERTIES_BY_DEFAULT)) {
            schema.put("additionalProperties", false);
        }

        if (Stream.of(schemaOptions).anyMatch((option) -> option == SchemaOption.UPPER_CASE_TYPE_VALUES)) {
            convertTypeValuesToUpperCase(schema);
        }
    }

    private static boolean isMethodParameterRequired(Method method, int index) {
        Parameter parameter = method.getParameters()[index];
        ToolParam toolParamAnnotation = parameter.getAnnotation(ToolParam.class);
        if (toolParamAnnotation != null) {
            return toolParamAnnotation.required();
        } else {
            JsonProperty propertyAnnotation = parameter.getAnnotation(JsonProperty.class);
            if (propertyAnnotation != null) {
                return propertyAnnotation.required();
            } else {
                Schema schemaAnnotation = parameter.getAnnotation(Schema.class);
                if (schemaAnnotation == null) {
                    Nullable nullableAnnotation = parameter.getAnnotation(Nullable.class);
                    return nullableAnnotation == null;
                } else {
                    return schemaAnnotation.requiredMode() == RequiredMode.REQUIRED
                            || schemaAnnotation.requiredMode() == RequiredMode.AUTO || schemaAnnotation.required();
                }
            }
        }
    }

    @Nullable
    private static String getMethodParameterDescription(Method method, int index) {
        Parameter parameter = method.getParameters()[index];
        ToolParam toolParamAnnotation = parameter.getAnnotation(ToolParam.class);
        if (toolParamAnnotation != null && StringUtils.hasText(toolParamAnnotation.description())) {
            return toolParamAnnotation.description();
        } else {
            JsonPropertyDescription jacksonAnnotation = parameter.getAnnotation(JsonPropertyDescription.class);
            if (jacksonAnnotation != null && StringUtils.hasText(jacksonAnnotation.value())) {
                return jacksonAnnotation.value();
            } else {
                Schema schemaAnnotation = parameter.getAnnotation(Schema.class);
                return schemaAnnotation != null && StringUtils.hasText(schemaAnnotation.description())
                        ? schemaAnnotation.description()
                        : null;
            }
        }
    }

    public static void convertTypeValuesToUpperCase(ObjectNode node) {
        if (node.isObject()) {
            node.fields().forEachRemaining((entry) -> {
                JsonNode value = entry.getValue();
                if (value.isObject()) {
                    convertTypeValuesToUpperCase((ObjectNode) value);
                } else if (value.isArray()) {
                    value.elements().forEachRemaining((element) -> {
                        if (element.isObject() || element.isArray()) {
                            convertTypeValuesToUpperCase((ObjectNode) element);
                        }
                    });
                } else if (value.isTextual() && entry.getKey().equals("type")) {
                    String oldValue = node.get("type").asText();
                    node.put("type", oldValue.toUpperCase());
                }
            });
        } else if (node.isArray()) {
            node.elements().forEachRemaining((element) -> {
                if (element.isObject() || element.isArray()) {
                    convertTypeValuesToUpperCase((ObjectNode) element);
                }
            });
        }
    }

    static {
        com.github.victools.jsonschema.generator.Module jacksonModule = new JacksonModule(
                JacksonOption.RESPECT_JSONPROPERTY_REQUIRED);
        com.github.victools.jsonschema.generator.Module openApiModule = new Swagger2Module();
        com.github.victools.jsonschema.generator.Module springAiSchemaModule = new JacksonModule();
        SchemaGeneratorConfigBuilder schemaGeneratorConfigBuilder = new SchemaGeneratorConfigBuilder(
                SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON)
                .with(jacksonModule)
                .with(openApiModule)
                .with(springAiSchemaModule)
                .with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .with(Option.PLAIN_DEFINITION_KEYS);
        SchemaGeneratorConfig typeSchemaGeneratorConfig = schemaGeneratorConfigBuilder.build();
        TYPE_SCHEMA_GENERATOR = new SchemaGenerator(typeSchemaGeneratorConfig);
        SUBTYPE_SCHEMA_GENERATOR = new SchemaGenerator(schemaGeneratorConfigBuilder.build());
    }

    public static enum SchemaOption {
        ALLOW_ADDITIONAL_PROPERTIES_BY_DEFAULT,
        UPPER_CASE_TYPE_VALUES;
    }
}
